from abc import ABC, abstractmethod
import random
import math as m

import cv2
import numpy as np
import skimage.transform as st
from auxmodels.utils.calc import xyxy2xywh, xywh2xyxy


class AugOperation(ABC):
    @abstractmethod
    def __call__(self, image, label):
        return image, label

    @staticmethod
    def _visualize_augment(image, label):
        image_plot = image.copy()
        for obj in label:
            p1 = (int(obj[0]), int(obj[1]))
            p2 = (int(obj[2]), int(obj[3]))
            cv2.rectangle(image_plot, p1, p2, (255, 0, 0), 2)
        cv2.imshow("vis", image_plot)
        cv2.waitKey()
        return


class RandomRotateAndCrop(AugOperation):
    """随机旋转和裁剪"""

    def __init__(
        self,
        min_size,
        max_angle=0,
        select_mode="center",
        wh_ratio=1.1,
        keep_ratio=0.2,
        p_rotate=0.5,
    ) -> None:
        super().__init__()
        assert select_mode in ["center", "all"]
        self.min_size = min_size
        self.max_angle = max_angle
        self.mode = select_mode
        self.wh_ratio = wh_ratio
        self.keep_ratio = keep_ratio
        self.p_rotate = p_rotate

    def __call__(self, image, label):
        picked_target = self._pick_center_target(image, label)
        image, label = self._rotate(image, label, picked_target)
        image, label = self._crop(image, label, picked_target)
        # self._assert_target_bbox_range(image, label)
        return image, label

    def _assert_target_bbox_range(self, image, label):
        h, w = image.shape[:2]
        assert (label[:, 0] < 0).sum() == 0, label[:, 0]
        assert (label[:, 1] < 0).sum() == 0, label[:, 1]
        assert (label[:, 2] > w).sum() == 0, (label[:, 2], label[:, 2] > w, w)
        assert (label[:, 3] > h).sum() == 0, (label[:, 3], label[:, 3] > h, h)
        return image, label

    def _pick_center_target(self, image, label):
        if len(label) > 0:
            h, w = image.shape[:2]
            if self.mode == "center":
                label = sorted(label, key=lambda x: abs(w / 2 - (x[2] + x[0]) / 2))
                picked_range = int((len(label) + 1) / 2)
            elif self.mode == "all":
                picked_range = len(label)
            picked_target = random.choice(label[:picked_range])
            return picked_target
        else:
            return

    def _rotate(self, image, label, center_target):
        if random.random() < self.p_rotate:
            h, w = image.shape[:2]
            if center_target is not None:
                # 旋转中心
                x0, y0, x1, y1 = center_target[:4]
                cx0, cy0 = int((x0 + x1) / 2), int((y0 + y1) / 2)
            else:
                cx0, cy0 = int(w / 2), int(h / 2)
                # 旋转角度
            rr = random.random()
            r_angle = (rr - 0.5) * 2 * self.max_angle
            # 图像旋转
            image = st.rotate(image, -r_angle, center=(cx0, cy0), preserve_range=True).astype(
                np.uint8
            )
            # 标注旋转
            if len(label) > 0:
                r_rad = r_angle / 180 * np.pi
                rmatrix = np.array([[m.cos(r_rad), m.sin(r_rad)], [-m.sin(r_rad), m.cos(r_rad)]])
                box_xywh = xyxy2xywh(label[:, :4])
                cxy = box_xywh[:, :2]
                cxy[:, 0] -= cx0
                cxy[:, 1] -= cy0
                cxy_new = np.matmul(cxy, rmatrix)
                cxy_new[:, 0] += cx0
                cxy_new[:, 1] += cy0
                box_xywh[:, :2] = cxy_new
                label[:, :4] = xywh2xyxy(box_xywh)
                maskx = np.logical_and(cxy_new[:, 0] > 0, cxy_new[:, 0] < w)
                masky = np.logical_and(cxy_new[:, 1] > 0, cxy_new[:, 1] < h)
                mask = np.logical_and(maskx, masky)
                label = label[mask]
                # self._visualize_augment(image, label)
                label = np.concatenate([label, np.ones((label.shape[0], 1))], axis=-1)
        else:
            if len(label) > 0:
                label = np.concatenate([label, np.zeros((label.shape[0], 1))], axis=-1)
        return image, label

    def _crop(self, image, label, center_target):
        if len(label) > 0:
            h, w = image.shape[:2]
            x0, y0, x1, y1 = center_target[:4]
            # w，h 的最小值
            w_min = max(self.min_size[1], abs(x1 - x0))
            h_min = max(self.min_size[0], abs(y1 - y0))
            # 目标中心
            cx = (x0 + x1) / 2
            cy = (y0 + y1) / 2
            # 剪裁区域的基准坐标
            x_base = max(0, cx - w_min / 2), min(w - 1, cx + w_min / 2)
            y_base = max(0, cy - h_min / 2), min(h - 1, cy + h_min / 2)
            # x方向外扩数量
            x_extend = random.random() * x_base[0], random.random() * (w - x_base[1])
            # x方向剪裁起始、结束坐标
            x_crop = int(x_base[0] - x_extend[0]), int(x_base[1] + x_extend[1])
            # 剪裁区域宽度
            w_crop = x_crop[1] - x_crop[0]
            # 剪裁区域高度
            h_crop = random.randrange(
                int(w_crop / w * h / self.wh_ratio), int(w_crop / w * h * self.wh_ratio)
            )
            # y方向外扩数量
            h_extend_total = h_crop - h_min
            h_extend = random.random() * h_extend_total
            y_crop = max(0, y_base[0] - h_extend)
            y_crop = int(y_crop), int(min(y_crop + h_crop, h - 1))
            # 删除剪裁区域外的targets
            label = np.array(label)
            crop_area = [x_crop[0], y_crop[0], x_crop[1], y_crop[1]]
            ratios = self._calc_crop_ratios(crop_area, label[:, :4])
            label_crop = label[ratios > self.keep_ratio]
            # 可视化检查
            # self._visualize_before_crop(image, label_crop, label, picked_target, crop_area)
            # 剪裁目标和图像
            image = image[y_crop[0] : y_crop[1], x_crop[0] : x_crop[1]]
            label_crop[:, 0] = (label_crop[:, 0] - x_crop[0]).clip(min=0)
            label_crop[:, 1] = (label_crop[:, 1] - y_crop[0]).clip(min=0)
            label_crop[:, 2] = (label_crop[:, 2] - x_crop[0]).clip(max=w_crop)
            label_crop[:, 3] = (label_crop[:, 3] - y_crop[0]).clip(max=y_crop[1] - y_crop[0])
            # 可视化检查
            # self._visualize_augment(image, label_crop)
            return image, label_crop
        else:
            return image, label

    # def _clip_targets(self, image, label):
    #     return image, label

    @staticmethod
    def _calc_crop_ratios(base, boxes):
        w1 = boxes[:, 2] - boxes[:, 0]
        h1 = boxes[:, 3] - boxes[:, 1]

        ix1 = boxes[:, 0].copy()
        ix1[ix1 < base[0]] = base[0]
        iy1 = boxes[:, 1].copy()
        iy1[iy1 < base[1]] = base[1]
        ix2 = boxes[:, 2].copy()
        ix2[ix2 > base[2]] = base[2]
        iy2 = boxes[:, 3].copy()
        iy2[iy2 > base[3]] = base[3]

        iw = (ix2 - ix1).clip(min=0.0)
        ih = (iy2 - iy1).clip(min=0.0)
        inters = iw * ih
        areas = w1 * h1

        ratios = inters / areas
        return ratios

    @staticmethod
    def _visualize_before_crop(image, label_crop, label_org, picked, crop_area):
        image_plot = image.copy()
        for obj in label_org:
            p1 = (int(obj[0]), int(obj[1]))
            p2 = (int(obj[2]), int(obj[3]))
            cv2.rectangle(image_plot, p1, p2, (255, 255, 0), 2)
        for obj in label_crop:
            p1 = (int(obj[0]), int(obj[1]))
            p2 = (int(obj[2]), int(obj[3]))
            cv2.rectangle(image_plot, p1, p2, (255, 0, 0), 2)
        p1 = (int(picked[0]), int(picked[1]))
        p2 = (int(picked[2]), int(picked[3]))
        cv2.rectangle(image_plot, p1, p2, (0, 255, 0), 2)
        p1 = (int(crop_area[0]), int(crop_area[1]))
        p2 = (int(crop_area[2]), int(crop_area[3]))
        cv2.rectangle(image_plot, p1, p2, (0, 0, 255), 2)
        cv2.imshow("crop", image_plot)
        cv2.waitKey()
        return


class RandomFlip(AugOperation):
    """随机翻转"""

    def __init__(self, horizontal: float = 0.5, vertical: float = 0) -> None:
        super().__init__()
        self.horizontal = horizontal
        self.vertical = vertical
        self.lut = [[None, 0], [1, -1]]

    def __call__(self, image, label):
        flip_mode_h = 1 if self.horizontal > random.random() else 0
        flip_mode_v = 1 if self.vertical > random.random() else 0
        flip_mode = self.lut[flip_mode_h][flip_mode_v]
        if flip_mode is None:
            return image, label
        else:
            h, w = image.shape[:2]
            label_temp = label.copy()
            image = cv2.flip(image, flip_mode)
            if len(label) > 0:
                if flip_mode == 0:  # vertical
                    label[:, 1] = h - label_temp[:, 1]
                    label[:, 3] = h - label_temp[:, 3]
                elif flip_mode == 1:  # horizontal
                    label[:, 0] = w - label_temp[:, 2]
                    label[:, 2] = w - label_temp[:, 0]
                elif flip_mode == -1:  # both
                    label[:, 0] = w - label_temp[:, 2]
                    label[:, 2] = w - label_temp[:, 0]
                    label[:, 1] = h - label_temp[:, 1]
                    label[:, 3] = h - label_temp[:, 3]
        # 可视化检查
        # self._visualize_augment(image, label)
        return image, label


class Resize(AugOperation):
    """尺度缩放"""

    def __init__(self, input_size) -> None:
        super().__init__()
        self.input_size = input_size

    def __call__(self, image, label):
        ih, iw = self.input_size
        h, w = image.shape[:2]
        rh = ih / h
        rw = iw / w
        image = cv2.resize(image, (iw, ih))
        if len(label) > 0:
            label[:, 0] *= rw
            label[:, 2] *= rw
            label[:, 1] *= rh
            label[:, 3] *= rh
        # self._visualize_augment(image, label)
        return image, label


class Normalize(AugOperation):
    """归一化"""

    def __init__(self, mean=128, scale=128) -> None:
        super().__init__()
        self.mean = mean
        self.scale = scale

    def __call__(self, image, label):
        image = (image - self.mean) / self.scale
        # self._visualize_augment(image, label)
        return image, label
